#  Copyright (c) 2010
#  The Regents of the University of Michigan
#  All Rights Reserved

#  Permission is granted to use, copy, create derivative works, and
#  redistribute this software and such derivative works for any purpose,
#  so long as the name of the University of Michigan is not used in
#  any advertising or publicity pertaining to the use or distribution
#  of this software without specific, written prior authorization. If
#  the above copyright notice or any other identification of the
#  University of Michigan is included in any copy of any portion of
#  this software, then the disclaimer below must also be included.

#  This software is provided as is, without representation or warranty
#  of any kind either express or implied, including without limitation
#  the implied warranties of merchantability, fitness for a particular
#  purpose, or noninfringement.  The Regents of the University of
#  Michigan shall not be liable for any damages, including special,
#  indirect, incidental, or consequential damages, with respect to any
#  claim arising out of or in connection with the use of the software,
#  even if it has been or is hereafter advised of the possibility of
#  such damages.
import collections
import logging
import socket
import struct

import gnutella

class SerializationError(Exception): pass

# Opcodes in Mojito.
# components/mojito/src/main/java/org/limewire/mojito/messages/DHTMessage.java
# Vuze name-a-likes come second in each line.
PING_REQUEST = ACT_REQUEST_PING = 1
PING_RESPONSE = ACT_REPLY_PING = 2
STORE_REQUEST = ACT_REQUEST_STORE = 3
STORE_RESPONSE = ACT_REPLY_STORE = 4
FIND_NODE_REQUEST = ACT_REQUEST_FIND_NODE = 5
FIND_NODE_RESPONSE = ACT_REPLY_FIND_NODE = 6
FIND_VALUE_REQUEST = ACT_REQUEST_FIND_VALUE = 7
FIND_VALUE_RESPONSE = ACT_REPLY_FIND_VALUE = 8
STATS_REQUEST = ACT_REQUEST_STATS = 9  # Deprecated per format wiki page.
STATS_RESPONSE = ACT_REPLY_STATS = 10  # Deprecated per format wiki page.

# Header flags.
FIREWALLED = 1
SHUTDOWN = 2

# General message structure.
GNUTELLA_HEADER = struct.Struct('!16sBH')
MOJITO_HEADER_TAIL = struct.Struct('!BBH')

CONTACT = struct.Struct('!4sH20sBIH') # address has 1-byte length prefix
CONTACT_ADDRESS = struct.Struct('!BIH')

# We decode the Gnutella header, but for now we're tossing it out and treating
# the Mojito header as the only useful header.
GnutellaHeader = collections.namedtuple('GnutellaHeader', ['guid',
                                                           'function',
                                                           'version',
                                                           'length'])
MojitoHeader = collections.namedtuple('MojitoHeader', ['opcode',
                                                       'contact',
                                                       'instance_id',
                                                       'flags',
                                                       'extended_length'])
                                                           
Contact = collections.namedtuple('Contact', ['vendor',
                                             'version',
                                             'id',
                                             'address'])

RequestPing = collections.namedtuple('RequestPing', ['header'])
ReplyPing = collections.namedtuple('ReplyPing', ['header',
                                                 'requester_address',
                                                 'dht_size'])

RequestFindNode = collections.namedtuple('RequestFindNode', ['header', 'id'])
ReplyFindNode = collections.namedtuple('ReplyFindNode', ['header',
                                                         'spoof_id',
                                                         'contacts'])
# Specific RPC bodies.
REPLY_PING_BODY = struct.Struct(CONTACT_ADDRESS.format + 'B')


def generate_gnutella_header():
  guid = gnutella.generate_guid()
  return GnutellaHeader(guid, ord('D'), gnutella.VERSION, 0)


def decode_fixup_ip_addr(addr):
  addr = struct.pack('!I', addr)
  return socket.inet_ntoa(addr)


def encode_fixup_ip_addr(addr):
  return struct.unpack('!I', socket.inet_aton(addr))[0]


def decode_big_integer(p):
  length = ord(p[0])
  return long(p[1:length + 1].encode('hex'), 16), p[length + 1:]


def encode_big_integer(i):
  raw = hex(i)[2:].decode('hex')
  return chr(len(raw)) + raw


def decode_contact(p):
  con = CONTACT  # Optimize global var lookup; this is a hot function.
  res = con.unpack(p[:con.size])
  if res[-3] != 4:
    raise SerializationError('Non-IPv4 address!')
  node_address = (decode_fixup_ip_addr(res[-2]), res[-1])
  return Contact(res[0], res[1], res[2], node_address), p[con.size:]


def encode_contact(c):
  ip = encode_fixup_ip_addr(c.address[0])
  p = CONTACT.pack(c.vendor, c.version, c.id, 4, ip, c.address[1])
  return p
  

def decode_gnutella_header(packet):
  res = GNUTELLA_HEADER.unpack(packet[:GNUTELLA_HEADER.size])
  packet = packet[GNUTELLA_HEADER.size:]
  length = struct.unpack('<I', packet[:4])

  return (GnutellaHeader(*(res + (length,))), packet[4:])


def encode_gnutella_header(hdr, new_len=None):
  if new_len is not None:
    hdr = GnutellaHeader(*(hdr[:-1] + (new_len,)))
  return (GNUTELLA_HEADER.pack(hdr.guid, hdr.function, hdr.version) +
          struct.pack('<I', hdr.length))


def decode_mojito_header(packet):
  opcode = ord(packet[0])
  packet = packet[1:]
  contact, packet = decode_contact(packet)
  return (MojitoHeader(opcode, contact, *MOJITO_HEADER_TAIL.unpack(packet[:MOJITO_HEADER_TAIL.size])),
          packet[MOJITO_HEADER_TAIL.size:])


def encode_mojito_header(hdr):
  p = chr(hdr.opcode)
  p += encode_contact(hdr.contact)
  p += MOJITO_HEADER_TAIL.pack(hdr.instance_id, hdr.flags, hdr.extended_length)
  return p


def finish_encoding_request(request, p):
  g_head = generate_gnutella_header()
  p = encode_mojito_header(request.header) + p
  p = encode_gnutella_header(g_head, len(p)) + p
  return p
  

def encode_ping_request(request):
  return finish_encoding_request(request, '')
  

def decode_ping_request(packet, header):
  return RequestPing(header)


def encode_find_node_request(request):
  return finish_encoding_request(request, request.id)


def decode_find_node_reply(packet, header):
  token, packet = decode_big_integer(packet)
  ncontacts = ord(packet[0])
  packet = packet[1:]

  contacts = []
  for _ in xrange(ncontacts):
    con, packet = decode_contact(packet)
    contacts.append(con)

  return ReplyFindNode(header, token, contacts)
    

def decode_ping_reply(packet, header):
  (_, ip, port) = CONTACT_ADDRESS.unpack(packet[:CONTACT_ADDRESS.size])
  ip = decode_fixup_ip_addr(ip)
  packet = packet[CONTACT_ADDRESS.size:]
  dht_size, packet = decode_big_integer(packet)

  return ReplyPing(header, (ip, port), dht_size)


def decode_packet(packet, handler, *args):
  header = None
  if not packet:
    raise SerializationError('Empty packet received!')

  try:
    gnutella_header, packet = decode_gnutella_header(packet)
    if gnutella_header.function != ord('D'):
      raise SerializationError('Non-DHT packet received; function %s' %
                               gnutella_header.function)
    header, packet = decode_mojito_header(packet)
    contact = header.contact
    if header.opcode == PING_REQUEST:
      handler.handle_ping_request(decode_ping_request(packet, header),
                                  contact,
                                  *args)
    elif header.opcode == PING_RESPONSE:
      handler.handle_ping_reply(decode_ping_reply(packet, header),
                                contact,
                                *args)
    elif header.opcode == FIND_NODE_RESPONSE:
      handler.handle_find_node_reply(decode_find_node_reply(packet, header),
                                     contact,
                                     *args)
    else:
      pass
#      print 'unknown opcode %d' % header.opcode

  except struct.error as e:
    if header is not None:
      logging.error('Opcode: %d' % header.opcode)
    raise SerializationError(e)
  except SerializationError as e:
    if header is not None:
      logging.debug('Opcode: %d' % header.opcode)
    raise

from cyserialization import decode_contact, decode_find_node_reply
from cyserialization import encode_find_node_request
